#include "sr2.h"

U32 const dithMask = 7;
U32 const mask23 = 0x3FFFF;
U32 const mask1  = 0x3F;
U32 const runThresh = 16;         /* encode 0 runs longer than this */

#define SHORTLITERALS 1

struct Srank
{
	Srank(std::vector<U8>& out)
	: out(out)
	{
		resetp();
		init();
		
	}

	U32 prev1, prev2, prev3;
	U32 Hx;

	void byte(U8 symbol)
	{
#if OLDHASH
		Hx = (prev1 | prev2 | prev3) & mask23 & ctxMask;
#endif

		U32 W = Contexts[Hx];         /* select context  */

		if((W & 0xFF) == symbol)           /* rank-0 match */
		{
#if NORUNS
			putBits(0, 1);
#else
			++runLength;
#endif
		}
		else if (((W >> 8) & 0xFF) == symbol) /* rank-1 match */
		{
			if (runLength > 0)
				emitRun();                         /* emit any pending run */
			putBits(6, 3);                   /* emit '110' */
			Contexts[Hx] =  (W & 0xFFFF0000) | ((W & 0xFF) << 8) | symbol;
		}
		else if (((W >> 16) & 0xFF) == symbol)  /* rank-2 match */
		{
			if (runLength > 0)
				emitRun();                         /* emit any pending run */
			putBits(14, 4);                  /* emit '1110' */
			Contexts[Hx] =  (W << 8) | symbol;
		}
		else                               /* no match */
		{
			if (runLength > 0)
				emitRun();                         /* emit any pending run */
			int SymIx = ChToIx[symbol];          /* get index in MTF table */

#if SHORTLITERALS
			if (SymIx < 32)                  /* near symbol */
#endif
			{
				/* emit '10' short Lit */
				putBits((2 << 5) + SymIx, 7);             /* emit index */
				//putBits(SymIx, 8);             /* and the index */
			}
#if SHORTLITERALS
			else                             /* remote symbol */
			{
				                                            /* emit '1111' */
				putBits((15 << 8) | SymIx, 12);             /* and the index */

				
			}
#endif
			Contexts[Hx] = (W << 8) | symbol;
		}

		pseudoMTF(symbol, ++dataBytes);

#if OLDHASH
		prev3 = (prev2 << 6);             /* move along hash codes */
		prev2 = ((prev1 & 0x3F) << 6);
		prev1 = symbol & mask1;           /* and get new one */
#else
		//Hx = ((Hx<<6) + (symbol & 0x3F)) & ctxMask;
		Hx = ((Hx*(5<<6)) + symbol) & ctxMask;
#endif
	}

	void putBits(U32 i, U32 bits)
	{
		if(out_bit_left >= bits)
		{
			// If bits is 32, then out_bit_left must also
			// be 32, and out_bits must be 0. Therefore
			// the shift distance doesn't matter.
			out_bits = (out_bits << bits) | i;
			out_bit_left -= bits;
		}
		else
		{
			// out_bit_left must be < 32 here, since
			// (out_bit_left < bits) && (bits <= 32).
			// left_over is in [1, 32], must use two shifts.
			U32 left_over = (bits - out_bit_left);
			U32 to_write = (out_bits << out_bit_left) | ((i >> 1) >> (left_over - 1));
			
			out_bit_left = (32-left_over);
			out_bits = (i << out_bit_left) >> out_bit_left;

			put_buf(to_write);
		}

#if 0
		// Immediate flush version:

		// out_bit_left in [1, 32]
		// bits in [1, 32]

		if(out_bit_left > bits)
		{
			// If bits is 32, we never get here, so (bits < 32)
			out_bits = (out_bits << bits) | i;
			out_bit_left -= bits;
		}
		else
		{
			// left_over >= 0 (since bits >= out_bit_left)
			// left_over is in [1, 32] - [1, 32] = [0, 31]
			// out_bit_left is in [1, 32], must use two shifts
			U32 left_over = (bits - out_bit_left);
			U32 to_write = ((out_bits << 1) << (out_bit_left - 1)) | (i >> left_over);
			
			out_bit_left = (32 - left_over);
			out_bits = (i & (1 << left_over)) - 1;

			put_buf(to_write);
		}
#endif
	}

	void resetp()
	{
		out_bits = 0;
		out_bit_left = 32;
	}

	void flush()
	{
		if (runLength > 0)
			emitRun();                         /* emit any pending run */
		//putBits(15, 4);
		//putBits( 0, 8);                      /* the EOF symbol */
		
		while(out_bit_left < 32)
		{
			out.push_back((out_bits >> 24) & 0xff);
			out_bits = (out_bits << 8);
			out_bit_left += 8;
		}

		resetp();
	}

	//#define SHORTRUNS 1

	void emitRun()
	{
	#if !SHORTLITERALS
		if (runLength <= 9)
		{
			putBits(0, runLength);    /* unary code for shorter runs */
		}
		else
		{
			int L;
			U32 mask;

			runLength -= 10;
			for (L = 0, mask = runLength; runLength; L ++)
				runLength >>= 1;
			putBits((15 << 5) | L, 9);          /* long Lit = 1111 */
			                                     /* bits to encode the run length */
			putBits(runLength, L);   /* and the run length itself */
		}
	#elif SHORTRUNS
	again:
		if (runLength <= 12)
		{
			putBits(0, runLength);    /* unary code for shorter runs */
		}
		else
		{
			runLength -= 13;

			if(runLength > 31)
			{
				putBits((15 << 8) | 31, 12);
				runLength -= 31;
				goto again;
			}

			putBits((15 << 8) | runLength, 12);
		}
	#else
		if (runLength <= runThresh)
		{
			putBits(0, runLength);    /* unary code for shorter runs */
		}
		else
		{
			int L;
			U32 mask;

			runLength -= runThresh;
			for (L = 0, mask = runLength; runLength; L ++)
				runLength >>= 1;
			putBits((15 << 8) | L, 12);          /* long Lit = 1111 */
			                                     /* bits to encode the run length */
			putBits(runLength, L);   /* and the run length itself */
		}
	#endif
		runLength = 0;             /* clear the length */
	}

	void put_buf(U32 w)
	{
		out.push_back((w >> 24) & 0xff);
		out.push_back((w >> 16) & 0xff);
		out.push_back((w >> 8) & 0xff);
		out.push_back(w & 0xff);
	}

	void init()
	{
		int i, ctxSize = 0x40000;

		Contexts.resize(ctxSize);

		for (i = 0; i < ctxSize; i++)
			Contexts[i] = 0x00652000;         /* initialise to  {NUL, ' ', 'e'} */

		ctxMask = (ctxSize-1);    /* get valid context index */
    
		/*compBytes = */
		dataBytes = 0;

		for (i = 0; i < 256; i++)
			ChToIx[i] = IxToCh[i] = i;        /* set up pseudo MTF table */

		prev1 = prev2 = prev3 = 0;
		runLength = 0;
		Hx = 0;
	}

	//#define OLDMTF 1

	void pseudoMTF(U8 currCh, U32 Dither)
    {
		int currIx, exchCh, exchIx;

		currIx = ChToIx[currCh];    /* position of current character */

	#if OLDMTF
		exchIx = currIx >> 1;       /* new posn of current character */
		exchIx ^= ((Dither) & dithMask);
	#else
		exchIx = Dither & 31;
	#endif
		exchCh = IxToCh[exchIx];

		ChToIx[currCh] = exchIx;
		IxToCh[exchIx] = currCh;
		ChToIx[exchCh] = currIx;
		IxToCh[currIx] = exchCh;
    }

	std::vector<U8>& out;

	int ChToIx[260];
	int IxToCh[260];

	std::vector<U32> Contexts;

	U32 ctxMask;

	U32 out_bits;
	U32 out_bit_left;

	U32 dataBytes;
	U32 runLength;
};

void srank_compress(std::vector<U8> const& in, std::vector<U8>& out)
{
	Srank c(out);

	for(std::size_t i = 0; i < in.size(); ++i)
		c.byte(in[i]);
	c.flush();
}
